# _*_ coding: utf-8 _*_
"""
Time:     2024/8/21
Author:   2051261
File:     db_operate.py
"""
import logging
from typing import Dict, Any

from sqlalchemy.ext.asyncio import async_sessionmaker
from sqlalchemy import text
from sqlalchemy.exc import IntegrityError
import pymysql
from mine.db.create_engine import ConnectionPoolConfig
from mine.db.create_engine import get_async_engine_from_config
from extractor.const import ConstKey
from extractor.const import Mark
from extractor.const import YmlDBConst


logger = logging.getLogger(__name__)


class DatabaseOperate:

    def __init__(self):
        self._session = {}
        self._name_type = {}

    async def add_session(self, conn_name, param):
        if conn_name in self._session:
            return
        engine = await get_async_engine_from_config(ConnectionPoolConfig(**param))
        self._session[conn_name] = async_sessionmaker(bind=engine, expire_on_commit=False)
        self._name_type[conn_name] = param[YmlDBConst.DATABASE_TYPE]

    async def get_db_type(self, conn_name):
        return self._name_type[conn_name]

    async def _ensure_connection(self, conn_name):
        if conn_name in self._session:
            session = self._session[conn_name]()
            connection = session.connection().connection
            try:
                connection.ping(reconnect=True)  # 如果连接断开，自动重连
            except Exception as e:
                logger.error(f"Error while pinging the connection: {str(e)}")
                raise e
            finally:
                session.close()

    async def _get_session(self, conn_name):
        if conn_name not in self._session:
            raise ValueError(f"No session found for connection '{conn_name}'")
        return self._session[conn_name]()

    def execute_query(self, conn_name, query, param=None):
        with self._get_session(conn_name) as session:
            try:
                result = session.execute(text(query), param)
                return result.fetchall()
            except Exception as e:
                logger.error(f"execute sql current error, sql:{query}, param:{param}")
                raise e

    async def fetch_one(self, conn_name, query, param=None):
        async with self._get_session(conn_name) as session:
            try:
                result = session.execute(text(query), param)
                row = result.fetchone()
                columns = result.keys()
                return dict(zip(columns, row)) if row else None
            except Exception as e:
                logger.error(f"execute sql current error, sql:{query}, param:{param}")
                raise e

    async def fetch_many(self, conn_name, query, param, batch_size=1000):
        async with self._get_session(conn_name) as session:
            try:
                result = session.execute(text(query), param)
                columns = result.keys()
                while True:
                    rows = result.fetchmany(batch_size)
                    if not rows:
                        break
                    for row in rows:
                        yield dict(zip(columns, row))
            except Exception as e:
                logger.error(f"execute sql current error, sql:{query}, param:{param}")
                raise e

    async def fetch_all(self, conn_name, query, param):
        async with self._get_session(conn_name) as session:
            try:
                result = session.execute(text(query), param)
                columns = result.keys()
                for row in result:
                    yield dict(zip(columns, row))
            except Exception as e:
                logger.error(f"execute sql current error, sql:{query}, param:{param}")
                raise e

    async def insert_one(self, conn_name, sql, param):
        async with self._get_session(conn_name) as session:
            try:
                result = session.execute(text(sql), param)
                session.commit()
                return {
                    ConstKey.AFFECTED_ROWS: result.rowcount,
                    ConstKey.LAST_ROW_ID: result.lastrowid,
                }
            except IntegrityError as e:
                session.rollback()
                if isinstance(e.orig, pymysql.err.IntegrityError) and e.orig.args[0] == 1062:
                    logger.warning(f"Unique key conflict occurred, sql:{sql}, param:{param}")
                    return Mark.UNIQUE_CONFLICT.value
                logger.error(f"occurred error:{e}, sql:{sql}, param:{param}")
                return Mark.INSERT_FAILED.value
            except Exception as e:
                logger.error(f"insert sql current error: {e}, sql:{sql}, param:{param}")
                session.rollback()
                return Mark.INSERT_FAILED.value

    async def update(self, conn_name: str, sql: str, param: Dict[str, Any]) -> Dict[str, Any]:
        async with self._get_session(conn_name) as session:
            try:
                result = session.execute(text(sql), param)
                session.commit()
                return {
                    ConstKey.AFFECTED_ROWS: result.rowcount,
                    ConstKey.LAST_ROW_ID: result.lastrowid,
                }
            except Exception as e:
                logger.error(f"update data current error: {e}, sql:{sql}, param:{param}")
                session.rollback()


db_operate = DatabaseOperate()
